# -*- coding: utf-8 -*-
"""
@Time    : 2024/9/2 19:07 
@Author  : ZhangShenao 
@File    : 2.利用条件边与循环流程实现工具调用Agent.py 
@Desc    : 利用条件边与循环流程实现工具调用Agent

在LangGraph中,Edge边定义了节点之间的通信逻辑和整个图架构的数据流向
边有以下4种类型:
1. 普通边: 固定从一个节点流向下一个节点
2. 条件边: 通过调用一个函数,来确定下一个需要跳转的节点
3. 入口边: 用户输入到达时首先调用的节点,即确定图结构的开始节点
4. 条件入口点: 调用一个函数来确定用户输入到达时首先调用的节点
            即通过函数来确定图结构的开始节点

一个节点可以拥有多个输出边
此时所有输出节点将在下一个步骤中并行执行

对于循环流程,LangGraph并没有封装单独的函数
可以通过add_edge函数,将两个节点串联到一起,并且指定条件边,来实现循环流程
需要注意的是:对于循环流程,一般都会有一个条件边用于跳出循环
否则LangGraph框架会检测到没有跳出循环的条件,应用程序会崩溃,并且极大消耗系统资源
"""
import json
import os
from typing import TypedDict, Annotated, Literal

import dotenv
from langchain_core.messages import ToolMessage
from langchain_core.runnables import RunnableConfig
from langchain_openai import ChatOpenAI
from langgraph.constants import END, START
from langgraph.graph import add_messages, StateGraph

from dalle_image_tool import create_dalle_image_tool
from gaode_weather_tool import GaodeWeatherTool
from google_serper_tool import create_google_serper_tool


# 定义图状态
class GraphState(TypedDict):
    """图状态"""

    messages: Annotated[list[str], add_messages]
    """
    状态字段——消息列表
    采用add_messages归纳函数,最终所有节点输出的消息都会聚合到该列表中
    """


# 加载环境变量
dotenv.load_dotenv()

# 创建工具列表
tools = [GaodeWeatherTool(), create_dalle_image_tool(), create_google_serper_tool()]

# 创建LLM,并绑定工具列表
llm = (ChatOpenAI(model="gpt-4o",
                  temperature=0,
                  openai_api_base=os.getenv("OPENAI_API_BASE"))
       .bind_tools(tools=tools))


# 定义LLM节点
def llm_node(state: GraphState, config: RunnableConfig) -> GraphState:
    # 从状态数据中获取消息列表
    messages = state["messages"]

    # 调用LLM模型,生成回复
    reply = llm.invoke(messages)

    # 返回状态数据,LangGraph框架会自动将每个节点返回的状态做聚合
    return {"messages": [reply]}


# 定义工具执行器节点
def tool_executor_node(state: GraphState, config: RunnableConfig) -> GraphState:
    """工具执行器节点"""

    # 构建工具映射字典
    tool_name_to_tool = {tool.name: tool for tool in tools}

    # 获取消息列表中的最后一条消息,即为工具调用消息
    tool_calls = state["messages"][-1].tool_calls

    # 遍历工具调用列表,实际调用工具,并将调用结果封装成ToolMessage,保存到消息列表中
    messages = []

    for tool_call in tool_calls:
        # 获取需要调用的工具
        tool_name = tool_call["name"]
        tool = tool_name_to_tool.get(tool_name)
        tool_id = tool_call["id"]

        # 调用工具,获取结果
        res = tool.invoke(tool_call["args"])

        # 将工具调用结果封装成ToolMessage,保存到消息列表中
        tool_msg = ToolMessage(content=json.dumps(res),
                               tool_call_id=tool_id,
                               name=tool_name)

        messages.append(tool_msg)

    # 将工具调用结果作为节点的执行状态返回
    return {"messages": messages}


# 定义节点的条件路由
def condition_route(state: GraphState, config: RunnableConfig) -> Literal["tool_executor", "__end__"]:
    """
    节点的条件路由
    根据状态数据,判断是跳转到工具调用节点,还是直接跳转到结束节点
    """

    # 获取状态中的最后一条消息
    last_msg = state["messages"][-1]

    # 如果消息中包含tool_calls字段,且长度>0
    # 则说明是工具调用消息,需要跳转到工具执行节点
    if hasattr(last_msg, "tool_calls") and len(last_msg.tool_calls) > 0:
        return "tool_executor"

    # 否则为普通文本消息,直接跳转到结束节点
    return END


# 初始化图架构,以GraphState作为状态数据
graph_builder = StateGraph(GraphState)

# 添加节点
graph_builder.add_node("llm", llm_node)
graph_builder.add_node("tool_executor", tool_executor_node)

# 添加普通边
graph_builder.add_edge(start_key=START, end_key="llm")
graph_builder.add_edge(start_key="tool_executor", end_key="llm")

# 添加条件边,并指定条件路由函数
graph_builder.add_conditional_edges(source="llm",  # 指定起始节点
                                    path=condition_route  # 指定路由函数,跟状态确定下一跳的节点
                                    )

# 将图架构编译为Runnable可运行组件
graph = graph_builder.compile()

# 定义初始状态
init_state = {"messages": [("human", "2024年奥运会男子乒乓球单打冠军是谁？")]}

# 调用图,获取执行结果
# 执行结果是一个GraphState状态
result = graph.invoke(init_state)
for msg in result["messages"]:
    if hasattr(msg, "tool_calls") and len(msg.tool_calls) > 0:
        print(f"工具调用参数: {msg.tool_calls}")

    print(f"消息内容: {msg.content}")
    print("---------------------------------")
